import pymysql
from langchain.agents import AgentExecutor, create_react_agent, create_tool_calling_agent
from langchain_openai import AzureChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.tools import tool

# Step 1: 设置 Azure OpenAI LLM
llm = AzureChatOpenAI(
    deployment_name="gpt-4o",
    openai_api_key='5aa0f77ee24c4b6e9910a131423bd02d',
    azure_endpoint='http://menshen.xdf.cn',
    openai_api_version='2024-02-15-preview',
    temperature=0,
)

# Step 2: 创建工具
@tool
def get_current_time(city: str) -> str:
    """返回某个城市当前的时间。"""
    from datetime import datetime
    return f"The current time in {city} is {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}."



@tool
def query_mysql_database(query: str) -> str:
    """
    Executes a SQL query on a MySQL database and returns the result.

    :param query: The SQL query to execute.
    :return: A string representation of the query result.
    """
    try:
        # Connect to the MySQL database
        connection = pymysql.connect(host="172.24.30.115", user="ourcrm_admin", password="ourcrm_admin", database="scrm_test")

        # Create a cursor to execute the query
        cursor = connection.cursor()
        cursor.execute(query)

        # Fetch the result
        result = cursor.fetchall()

        # Close the connection
        cursor.close()
        connection.close()

        # Return the result as a string
        return str(result)

    except Exception as e:
        return f"Error executing query: {str(e)}"

tools = [get_current_time, query_mysql_database]

# Step 3: 自定义 Prompt
template = """你是一个智能体，擅长使用工具解决问题。
你可以使用以下工具：
{tools}

请使用如下格式：

Question: {input}
{agent_scratchpad}

Thought: 你在思考什么？
Action: 工具名称（{tool_names} 中的一个）
Action Input: 工具的输入
Observation: 工具返回的结果
...（可以重复 Thought/Action/Action Input/Observation 多次）
Thought: 我知道最终答案了
Final Answer: 最终答案是什么

现在开始！

Question: {input}
{agent_scratchpad}
"""

prompt = PromptTemplate.from_template(template)

# Step 4: 创建 agent 和 executor
# agent = create_react_agent(llm=llm, tools=tools, prompt=prompt)
agent = create_tool_calling_agent(llm=llm, tools=tools, prompt=prompt)
executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# Step 5: 调用
if __name__ == "__main__":
    result = executor.invoke({"input": "请告诉我现在上海几点了？之后执行一下这个sql查询：select * from crm_student limit 1，获取其中的学员的Id，根据学员Id查询对应的商机"})
    print("最终结果：", result)